import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
import matplotlib.pyplot as plt
import numpy as np
from sklearn.metrics import classification_report, confusion_matrix

# 1. 加载数据
iris = load_iris()
X, y = iris.data, iris.target

# 2. 数据标准化 + reshape 为 (batch, channel, height, width)
scaler = StandardScaler()
X = scaler.fit_transform(X)
X = X.reshape(-1, 1, 2, 2)  # CNN输入：(batch, channel, H, W)

# 转为张量
X = torch.tensor(X, dtype=torch.float32)
y = torch.tensor(y, dtype=torch.long)

# 3. 划分训练测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)

# 4. 定义 CNN 网络结构
class IrisCNN(nn.Module):
    def __init__(self):
        super(IrisCNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 8, kernel_size=1)
        self.relu = nn.ReLU()
        self.flatten = nn.Flatten()
        self.fc1 = nn.Linear(8 * 2 * 2, 16)
        self.fc2 = nn.Linear(16, 3)  # 三分类

    def forward(self, x):
        x = self.relu(self.conv1(x))
        x = self.flatten(x)
        x = self.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 实例化模型
model = IrisCNN()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.01)

# 5. 训练模型
epochs = 100
train_acc_history = []

for epoch in range(epochs):
    model.train()
    optimizer.zero_grad()
    outputs = model(X_train)
    loss = criterion(outputs, y_train)
    loss.backward()
    optimizer.step()

    # 计算训练准确率
    _, preds = torch.max(outputs, 1)
    acc = (preds == y_train).float().mean()
    train_acc_history.append(acc.item())

# 6. 测试模型
model.eval()
with torch.no_grad():
    test_outputs = model(X_test)
    _, test_preds = torch.max(test_outputs, 1)
    test_acc = (test_preds == y_test).float().mean()
    print("\n=== 测试集准确率: {:.2f}% ===".format(test_acc.item() * 100))
    print("分类报告：\n", classification_report(y_test.numpy(), test_preds.numpy()))

# 7. 可视化训练准确率
plt.figure(figsize=(8, 4))
plt.plot(train_acc_history, label='Train Accuracy')
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.title('Training Accuracy Over Epochs')
plt.legend()
plt.grid(True)
plt.show()

# 8. 可视化预测 vs 实际
plt.figure(figsize=(8, 4))
plt.plot(test_preds.numpy(), 'ro-', label='Predicted')
plt.plot(y_test.numpy(), 'go--', label='Actual')
plt.title("Predicted vs Actual on Test Set")
plt.xlabel("Sample Index")
plt.ylabel("Class")
plt.legend()
plt.grid(True)
plt.show()
